# index_manager.py
# -*- coding: utf-8 -*-

import os
import logging
import time
import pickle
import faiss
import numpy as np
from typing import Dict, Any, List, Tuple, Set, Optional
from collections import defaultdict, Counter

# 导入项目模块
import config
from plagiarism_core.core_engine import ProcessedDocData # 导入预处理数据结构类型
from datasketch import MinHash, MinHashLSH


logger = logging.getLogger(__name__)

class IndexManager:
    """
    管理查重项目中使用的各种索引。
    """

    def __init__(self, corpus_data: Dict[str, ProcessedDocData]):
        """
        初始化索引管理器。
        """
        logger.info("开始初始化 IndexManager...")
        self.corpus_data = corpus_data
        self.index_dir = config.INDEX_DIR # 索引文件存储目录
        os.makedirs(self.index_dir, exist_ok=True) # 确保目录存在

        # --- 初始化 MinHash LSH ---
        self.lsh_index: Optional[MinHashLSH] = None
        self.corpus_minhashes: Dict[str, MinHash] = {} # 存储每个语料的 MinHash 签名
        self.minhash_enabled: bool = False
        if config.USE_MINHASH_PRESCREENING:
            self._load_or_build_minhash_lsh()
        else:
            logger.info("配置中已禁用 MinHash LSH。")

        # --- 初始化倒排索引 ---
        self.inverted_index: Optional[Dict[str, List[Tuple[str, int]]]] = None # {ngram: [(doc_id, pos), ...]}
        self.inverted_index_enabled: bool = False
        if config.USE_INVERTED_INDEX:
            self._load_or_build_inverted_index()
        else:
            logger.info("配置中已禁用倒排索引。")

        # --- 初始化语义向量索引 (Faiss) ---
        self.semantic_index: Optional[faiss.Index] = None # Faiss 索引对象

        # 映射 Faiss 索引中的向量 ID -> (语料库文档 ID, 块在文档中的索引, 块起始位置, 块结束位置)
        self.semantic_id_map: List[Tuple[str, int, int, int]] = []
        self.semantic_index_enabled: bool = False

        # 语义索引的构建可能依赖语义模型，加载/构建逻辑放在这里
        if config.USE_SEMANTIC_INDEX:
            self._load_semantic_index() # 先尝试加载，如果加载失败或不存在，需要手动调用构建方法
            if not self.semantic_index_enabled:
                 logger.warning("语义索引未加载成功。如果需要使用，请确保已构建或检查索引文件。")
        else:
            logger.info("配置中已禁用语义向量索引。")

        logger.info("IndexManager 初始化完成。")

    # === MinHash LSH 相关方法 ===

    def _get_minhash_lsh_path(self) -> str:
        """获取 MinHash LSH 索引文件的标准路径。"""
        return os.path.join(self.index_dir, "minhash_lsh.pkl")

    def _load_or_build_minhash_lsh(self):
        """加载或构建 MinHash LSH 索引。"""
        if not self.corpus_data:
             logger.warning("无法构建 MinHash LSH 索引：语料库数据为空。")
             return

        index_path = self._get_minhash_lsh_path()

        # 尝试加载现有索引
        if os.path.exists(index_path):
            logger.info(f"发现 MinHash LSH 索引文件，尝试加载: {index_path}")
            try:
                with open(index_path, 'rb') as f:
                    loaded_data = pickle.load(f)
                # 验证加载的数据结构和内容
                if isinstance(loaded_data, tuple) and len(loaded_data) == 2:
                    lsh, minhashes = loaded_data
                    if isinstance(lsh, MinHashLSH) and isinstance(minhashes, dict):
                        # 检查索引是否与当前语料库大致匹配 (简单检查数量)
                        # 更严格的检查可以比对 keys 或使用版本号/校验和
                        if len(minhashes) == len(self.corpus_data):
                            self.lsh_index = lsh
                            self.corpus_minhashes = minhashes
                            self.minhash_enabled = True
                            logger.info(f"MinHash LSH 索引加载成功 ({len(self.corpus_minhashes)} 个签名)。")
                            return # 加载成功，直接返回
                        else:
                            logger.warning(f"加载的 MinHash 签名数量 ({len(minhashes)}) 与当前语料库文档数 ({len(self.corpus_data)}) 不匹配。将重新构建索引。")
                    else:
                        logger.warning("加载的 MinHash LSH 数据类型不正确。将重新构建索引。")
                else:
                    logger.warning("加载的 MinHash LSH 数据格式不正确。将重新构建索引。")
            except Exception as e:
                logger.exception(f"加载 MinHash LSH 索引失败: {e}。将重新构建索引。")
        else:
             logger.info("未找到 MinHash LSH 索引文件，将开始构建新索引。")

        # --- 如果加载失败或文件不存在，则构建索引 ---
        self._build_and_save_minhash_lsh(index_path)

    def _build_and_save_minhash_lsh(self, save_path: str):
        """构建 MinHash LSH 索引并保存到文件。"""
        logger.info(f"开始构建 MinHash LSH 索引 (保存路径: {save_path})...")
        start_time = time.time()

        # 初始化 LSH 对象
        lsh = MinHashLSH(threshold=config.MINHASH_THRESHOLD,
                         num_perm=config.MINHASH_NUM_PERM)
        minhashes = {} # 存储计算出的 MinHash 签名
        processed_count = 0
        skipped_count = 0

        # 遍历语料库数据计算 MinHash 并插入 LSH
        for corpus_id, data in self.corpus_data.items():
            text = data.get("cleaned_text", "")
            # 可以在这里添加长度过滤，或者在 _calculate_minhash 内部处理
            if text and len(text) >= config.MINHASH_INPUT_MIN_CHARS_THRESHOLD:
                minhash_sig = self._calculate_minhash(text) # 调用辅助函数计算
                if minhash_sig:
                    minhashes[corpus_id] = minhash_sig
                    try:
                        # LSH 插入需要 bytes 类型的 key
                        lsh.insert(corpus_id.encode('utf-8'), minhash_sig)
                        processed_count += 1
                    except Exception as insert_err:
                        logger.error(f"将文档 '{corpus_id}' 的 MinHash 插入 LSH 时出错: {insert_err}", exc_info=True)
                        # 从 minhashes 中移除，避免数据不一致
                        minhashes.pop(corpus_id, None)
                else:
                    logger.debug(f"文档 '{corpus_id}' 未能生成有效的 MinHash 签名（可能过短或处理失败）。")
                    skipped_count += 1
            else:
                 logger.debug(f"跳过文档 '{corpus_id}' 的 MinHash 计算，因文本为空或长度不足 ({len(text)} < {config.MINHASH_INPUT_MIN_CHARS_THRESHOLD})。")
                 skipped_count += 1

            # 打印进度
            total_processed = processed_count + skipped_count
            if total_processed > 0 and (total_processed % 200 == 0 or total_processed == len(self.corpus_data)):
                logger.info(f"  已处理 {total_processed}/{len(self.corpus_data)} 个文档用于 MinHash LSH (成功: {processed_count}, 跳过: {skipped_count})...")

        if processed_count == 0:
            logger.warning("未能成功为任何文档生成 MinHash 签名并添加到 LSH 索引。索引将为空，MinHash 功能将禁用。")
            self.lsh_index = lsh # 保存空索引
            self.corpus_minhashes = minhashes
            self.minhash_enabled = False
        else:
            logger.info(f"MinHash 计算和 LSH 插入完成。共处理 {processed_count} 个文档。")
            self.lsh_index = lsh
            self.corpus_minhashes = minhashes
            self.minhash_enabled = True # 只有成功处理了文档才启用

        # --- 保存索引和签名到文件 ---
        try:
            os.makedirs(os.path.dirname(save_path), exist_ok=True)
            with open(save_path, 'wb') as f:
                pickle.dump((self.lsh_index, self.corpus_minhashes), f)
            logger.info(f"MinHash LSH 索引和签名已成功保存到: {save_path}")
        except Exception as e:
            logger.exception(f"保存 MinHash LSH 索引文件失败: {e}")

        build_duration = time.time() - start_time
        logger.info(f"MinHash LSH 索引构建完成，耗时 {build_duration:.2f} 秒。状态: {'启用' if self.minhash_enabled else '禁用'}")

    # --- MinHash 计算辅助方法 ---
    @staticmethod
    def _create_shingles(text: str, shingle_size: int) -> set:
        """为 MinHash 生成字符级 Shingles (n-grams)。"""
        shingles = set()
        if not text or len(text) < shingle_size:
            return shingles
        # 移除文本中的所有空白字符，避免影响 shingle
        from plagiarism_core.preprocessor import RE_HORIZONTAL_WHITESPACE # 临时导入
        text_no_space = RE_HORIZONTAL_WHITESPACE.sub('', text.replace('\n', '').replace('\r', ''))
        if len(text_no_space) < shingle_size:
             return shingles
        # 生成 shingles 并编码为 bytes
        for i in range(len(text_no_space) - shingle_size + 1):
            shingle = text_no_space[i:i + shingle_size]
            shingles.add(shingle.encode('utf-8'))
        return shingles

    @staticmethod
    def _calculate_minhash(text: str) -> Optional[MinHash]:
        """计算给定文本的 MinHash 签名。"""
        if not text: return None

        try:
            shingle_size = config.MINHASH_SHINGLE_SIZE
            num_perm = config.MINHASH_NUM_PERM
            shingles = IndexManager._create_shingles(text, shingle_size)
            if not shingles:
                # logger.debug("文本未能生成 shingles (可能过短或只有空白)。")
                return None
            # 创建 MinHash 对象
            m = MinHash(num_perm=num_perm)
            # 更新签名
            for s in shingles:
                m.update(s)
            return m
        except Exception as e:
            logger.error(f"计算 MinHash 时发生错误: {e}", exc_info=False) # 减少日志冗余
            return None

    # === 倒排索引相关方法 ===

    def _get_inverted_index_path(self) -> str:
        """获取倒排索引文件的标准路径。"""
        return os.path.join(self.index_dir, "inverted_index.pkl")

    def _load_or_build_inverted_index(self):
        """加载或构建倒排索引。"""
        if not self.corpus_data:
             logger.warning("无法构建倒排索引：语料库数据为空。")
             return

        index_path = self._get_inverted_index_path()

        # 尝试加载
        if os.path.exists(index_path):
            logger.info(f"发现倒排索引文件，尝试加载: {index_path}")
            try:
                with open(index_path, 'rb') as f:
                    loaded_index = pickle.load(f)
                if isinstance(loaded_index, dict):
                     # TODO: 添加更严格的验证，例如检查索引版本或部分词项
                     self.inverted_index = loaded_index
                     self.inverted_index_enabled = True
                     logger.info(f"倒排索引加载成功 ({len(self.inverted_index)} 个词项)。")
                     return
                else:
                     logger.warning("加载的倒排索引数据类型不正确。将重新构建。")
            except Exception as e:
                logger.exception(f"加载倒排索引失败: {e}。将重新构建。")
        else:
             logger.info("未找到倒排索引文件，将开始构建新索引。")

        # --- 构建索引 ---
        self._build_and_save_inverted_index(index_path)

    def _build_and_save_inverted_index(self, save_path: str):
        """构建基于 N-gram 的倒排索引并保存。"""
        logger.info(f"开始构建倒排索引 (保存路径: {save_path})...")
        start_time = time.time()

        inverted_index_temp = defaultdict(list) # 使用临时变量构建
        ngram_size = config.INVERTED_INDEX_NGRAM_SIZE
        processed_docs = 0
        total_postings = 0

        # 遍历语料库生成 N-grams 并添加到索引
        for doc_id, data in self.corpus_data.items():
            text = data.get("cleaned_text", "")
            if not text or len(text) < ngram_size:
                continue # 跳过过短或空的文本

            processed_docs += 1
            # --- 生成 N-grams ---
            for i in range(len(text) - ngram_size + 1):
                ngram = text[i : i + ngram_size]
                # 过滤掉包含换行符或完全是空白的 N-gram
                # 也可以添加其他过滤条件，例如是否包含标点符号
                if '\n' not in ngram and ngram.strip():
                    # 存储 (文档ID, N-gram 在文档中的起始位置)
                    inverted_index_temp[ngram].append((doc_id, i))
                    total_postings += 1

            # 打印进度
            if processed_docs % 200 == 0 or processed_docs == len(self.corpus_data):
                logger.info(f"  已处理 {processed_docs}/{len(self.corpus_data)} 个文档用于倒排索引...")

        if not inverted_index_temp:
            logger.warning("未能从语料库生成任何有效的 N-gram 用于倒排索引。索引将为空。")
            self.inverted_index = {}
            self.inverted_index_enabled = False
            # 即使为空，也保存一个空文件，避免下次重复构建
            try:
                with open(save_path, 'wb') as f: pickle.dump({}, f)
            except Exception as e: logger.exception(f"保存空的倒排索引文件失败: {e}")
            return

        logger.info(f"初步生成了 {len(inverted_index_temp)} 个唯一 N-gram 索引条目，包含 {total_postings} 个倒排记录。")

        # --- 应用文档频率 (DF) 过滤 ---
        min_df = config.INVERTED_INDEX_MIN_DF
        max_df_ratio = config.INVERTED_INDEX_MAX_DF_RATIO
        num_docs = len(self.corpus_data) # 使用加载的文档总数
        # 计算最大文档频率绝对值
        max_df = int(num_docs * max_df_ratio) if max_df_ratio < 1.0 else num_docs
        # 确保 max_df 不小于 min_df
        if max_df < min_df:
             logger.warning(f"计算得到的倒排索引 max_df ({max_df}) 小于 min_df ({min_df})，已将 max_df 设为文档总数 {num_docs}。")
             max_df = num_docs

        if min_df > 1 or max_df < num_docs:
            logger.info(f"开始根据文档频率过滤倒排索引 (min_df={min_df}, max_df={max_df}, max_df_ratio={max_df_ratio:.2f})...")
            filtered_index = {}
            original_ngram_count = len(inverted_index_temp)
            filtered_out_count = 0
            for ngram, postings in inverted_index_temp.items():
                # 计算包含该 N-gram 的独立文档数量 (Document Frequency)
                doc_freq = len(set(doc_id for doc_id, _ in postings))
                # 应用过滤条件
                if min_df <= doc_freq <= max_df:
                    filtered_index[ngram] = postings
                else:
                     filtered_out_count += 1
                     # logger.debug(f"  过滤 N-gram '{ngram}' (DF={doc_freq})")

            if filtered_out_count > 0:
                inverted_index_temp = filtered_index # 使用过滤后的索引
                logger.info(f"倒排索引过滤完成，移除了 {filtered_out_count} 个 N-gram (占比 {(filtered_out_count/original_ngram_count)*100:.2f}%)，"
                            f"剩余 {len(inverted_index_temp)} 个。")
            else:
                logger.info("倒排索引过滤未移除任何 N-gram。")
        # -------------------------

        self.inverted_index = dict(inverted_index_temp) # 存储最终索引
        self.inverted_index_enabled = True

        # --- 保存索引到文件 ---
        try:
            os.makedirs(os.path.dirname(save_path), exist_ok=True)
            with open(save_path, 'wb') as f:
                pickle.dump(self.inverted_index, f)
            logger.info(f"倒排索引已成功保存到: {save_path}")
        except Exception as e:
            logger.exception(f"保存倒排索引文件失败: {e}")
            # 保存失败不应影响内存中的对象状态

        build_duration = time.time() - start_time
        logger.info(f"倒排索引构建完成，耗时 {build_duration:.2f} 秒。状态: {'启用' if self.inverted_index_enabled else '禁用'}")

    # === 语义向量索引 (Faiss) 相关方法 ===

    def _get_semantic_index_paths(self) -> Tuple[str, str]:
        """获取语义索引文件和映射文件的标准路径。"""
        index_path = os.path.join(self.index_dir, "semantic_index.faiss")
        map_path = os.path.join(self.index_dir, "semantic_mapping.pkl")
        return index_path, map_path

    def _load_semantic_index(self):
        """尝试加载预先构建的 Faiss 索引和 ID 映射。"""
        index_path, map_path = self._get_semantic_index_paths()

        if os.path.exists(index_path) and os.path.exists(map_path):
            logger.info(f"发现语义索引文件，尝试加载: {index_path} 和 {map_path}")
            try:
                # 加载 Faiss 索引
                semantic_index = faiss.read_index(index_path)
                logger.info(f"  Faiss 索引加载成功: 维度={semantic_index.d}, 向量总数={semantic_index.ntotal}")

                # 加载 ID 映射
                with open(map_path, 'rb') as f_map:
                    semantic_id_map = pickle.load(f_map)
                logger.info(f"  ID 映射加载成功，包含 {len(semantic_id_map)} 条记录。")

                # --- 验证索引和映射的一致性 ---
                if semantic_index.ntotal == len(semantic_id_map):
                    # 检查映射中的文档 ID 是否存在于当前语料库中
                    # （允许索引包含旧文档，但不能包含未知文档）
                    mapped_ids = {item[0] for item in semantic_id_map} # 提取所有文档 ID
                    current_ids = set(self.corpus_data.keys())
                    unknown_ids = mapped_ids - current_ids
                    if not unknown_ids: # 如果所有映射的 ID 都在当前语料库中
                        self.semantic_index = semantic_index
                        self.semantic_id_map = semantic_id_map
                        self.semantic_index_enabled = True
                        logger.info("语义索引和映射加载并验证成功。")
                        return # 加载成功
                    else:
                         logger.warning(f"语义索引映射包含 {len(unknown_ids)} 个当前语料库中不存在的文档 ID (例如: {list(unknown_ids)[:5]})。索引可能已过时，建议重新构建。")
                         # 仍然加载，但标记为可能不准确或禁用？当前选择是继续加载
                         self.semantic_index = semantic_index
                         self.semantic_id_map = semantic_id_map
                         self.semantic_index_enabled = True # 暂时仍启用
                         # return
                else:
                    logger.error(f"语义索引大小 ({semantic_index.ntotal}) 与 ID 映射大小 ({len(semantic_id_map)}) 不匹配！索引文件已损坏或不一致。")
                    # 不加载不一致的数据

            except FileNotFoundError:
                 logger.error(f"尝试加载语义索引时文件未找到 (即使之前检查存在): {index_path} 或 {map_path}")
            except Exception as e:
                logger.exception(f"加载语义索引或映射文件时发生错误: {e}")
        else:
            logger.info("未找到预构建的语义索引文件 (semantic_index.faiss 和/或 semantic_mapping.pkl)。")
            # 不执行任何操作，semantic_index_enabled 保持 False

    def build_semantic_index(self, semantic_model: Any):
        """
        构建 Faiss 语义向量索引并保存。
        需要一个已初始化的 Sentence Transformer 模型实例来计算嵌入。
        """
        if not semantic_model:
             logger.error("无法构建语义索引：未提供有效的语义模型实例。")
             return False
        if not self.corpus_data:
             logger.warning("无法构建语义索引：语料库数据为空。")
             return False

        index_path, map_path = self._get_semantic_index_paths()
        logger.info(f"开始构建语义向量索引 (保存路径: {index_path}, {map_path})...")
        start_time = time.time()

        all_vectors_list = [] # 存储每个文档的嵌入 numpy 数组
        id_map: List[Tuple[str, int, int, int]] = [] # (doc_id, chunk_idx_in_doc, start, end)
        try:
             vector_dim = semantic_model.get_sentence_embedding_dimension()
        except Exception as e:
             logger.error(f"无法获取语义模型的向量维度: {e}")
             return False
        processed_chunks_count = 0
        processed_docs_count = 0
        model_device = getattr(semantic_model, 'device', 'cpu') # 尝试获取模型设备

        logger.info(f"向量维度: {vector_dim}。模型运行设备: {model_device}")
        logger.info(f"将使用 '{config.SEMANTIC_CHUNK_TYPE}' 作为索引的基本单元。")

        # --- 遍历语料库，获取块并计算嵌入 ---
        for corpus_id, data in self.corpus_data.items():
            processed_docs_count += 1
            cleaned_text = data.get("cleaned_text", "")
            if not cleaned_text:
                logger.debug(f"跳过空文档 '{corpus_id}' 的语义索引构建。")
                continue

            # 根据配置从预处理结果中获取分块
            chunks_with_pos: List[Tuple[str, int, int]] = data.get(
                "sentences_pos" if config.SEMANTIC_CHUNK_TYPE == 'sentence' else "paragraphs_pos", []
            )
            if not chunks_with_pos:
                 logger.debug(f"文档 '{corpus_id}' 没有找到有效的 '{config.SEMANTIC_CHUNK_TYPE}' 块用于索引。")
                 continue

            chunk_texts = [c[0] for c in chunks_with_pos]
            logger.debug(f"文档 '{corpus_id}': 正在编码 {len(chunk_texts)} 个块...")

            try:
                # --- 批量计算嵌入 ---
                # 使用 convert_to_numpy=True 直接获取 numpy 数组
                embeddings = semantic_model.encode(
                    chunk_texts,
                    convert_to_numpy=True,
                    show_progress_bar=False, # 避免过多日志
                    device=model_device,
                    batch_size=config.SEMANTIC_ENCODE_BATCH_SIZE
                )
                # -------------------
                if not isinstance(embeddings, np.ndarray) or embeddings.ndim != 2:
                     logger.error(f"文档 '{corpus_id}' 的嵌入计算结果类型或维度不正确。跳过此文档。")
                     continue
                if embeddings.shape[0] != len(chunk_texts):
                    logger.error(f"文档 '{corpus_id}' 的嵌入向量数量 ({embeddings.shape[0]}) 与块数量 ({len(chunk_texts)}) 不匹配。跳过此文档。")
                    continue
                if embeddings.shape[1] != vector_dim:
                    logger.error(f"文档 '{corpus_id}' 的嵌入向量维度 ({embeddings.shape[1]}) 与预期 ({vector_dim}) 不符。跳过此文档。")
                    continue

                # 检查嵌入中是否有 NaN 或 Inf (可能导致 Faiss 出错)
                if np.isnan(embeddings).any() or np.isinf(embeddings).any():
                     logger.warning(f"文档 '{corpus_id}' 的嵌入向量中包含 NaN 或 Inf 值！将尝试替换为零向量，但这可能影响准确性。")
                     embeddings = np.nan_to_num(embeddings, nan=0.0, posinf=0.0, neginf=0.0) # 替换为 0

                all_vectors_list.append(embeddings.astype('float32')) # 确保是 float32

                # --- 创建 ID 映射 ---
                for i, (_, start_char, end_char) in enumerate(chunks_with_pos):
                    # 映射存储: (文档ID, 块在文档内的索引, 块起始位置, 块结束位置)
                    id_map.append((corpus_id, i, start_char, end_char))
                # --------------------

                processed_chunks_count += len(chunk_texts)
                # 打印进度
                if processed_docs_count % 50 == 0 or processed_docs_count == len(self.corpus_data):
                    logger.info(f"  已编码 {processed_docs_count}/{len(self.corpus_data)} 个文档的块...")

            except Exception as e:
                logger.error(f"编码文档 '{corpus_id}' 的块时发生错误: {e}", exc_info=True)
                # 如果单个文档失败，可以选择跳过它继续构建，或者中止整个构建过程

        if not all_vectors_list:
            logger.error("未能为任何文档生成有效的嵌入向量。无法构建语义索引。")
            return False

        logger.info(f"语料库嵌入计算完成。共生成 {len(id_map)} 个向量 (来自 {processed_chunks_count} 个块)。")

        # --- 构建 Faiss 索引 ---
        try:
            # 将列表中的 numpy arrays 合并为一个大的 numpy array
            vectors_np = np.vstack(all_vectors_list)
            if vectors_np.shape[0] != len(id_map):
                logger.error(f"最终向量数量 ({vectors_np.shape[0]}) 与 ID 映射数量 ({len(id_map)}) 不匹配！索引构建失败。")
                return False

            # --- 向量归一化 (对于内积或余弦相似度是推荐的) ---
            logger.info("正在对向量进行 L2 归一化...")
            faiss.normalize_L2(vectors_np)
            logger.info("向量归一化完成。")

            # 创建 Faiss 索引
            index_type = config.FAISS_INDEX_TYPE
            # 对于内积 (METRIC_INNER_PRODUCT)，向量归一化后，内积等价于余弦相似度
            metric = faiss.METRIC_INNER_PRODUCT
            logger.info(f"开始构建 Faiss 索引，类型: {index_type}，度量: 内积 (Inner Product)")
            index = faiss.index_factory(vector_dim, index_type, metric)

            # --- 训练索引 (如果需要) ---
            if not index.is_trained:
                train_size = min(vectors_np.shape[0], config.FAISS_TRAIN_SIZE)
                if vectors_np.shape[0] > index.ntotal: # 检查是否有足够的数据训练
                     logger.info(f"开始训练 Faiss 索引 (使用 {train_size} 个向量)...")
                     # 如果数据量大，随机抽样训练；否则用全部数据
                     if train_size < vectors_np.shape[0]:
                          np.random.seed(42) # for reproducibility
                          random_indices = np.random.choice(vectors_np.shape[0], train_size, replace=False)
                          index.train(vectors_np[random_indices])
                     else:
                          index.train(vectors_np)
                     logger.info("Faiss 索引训练完成。")
                else:
                      logger.warning("没有足够的数据或索引类型不需要训练。")

            # --- 添加向量到索引 ---
            # 使用 add_with_ids 将向量 ID (0 到 N-1) 与向量关联起来
            if vectors_np.shape[0] > 0:
                logger.info(f"开始向 Faiss 索引添加 {vectors_np.shape[0]} 个向量...")
                ids_np = np.arange(vectors_np.shape[0]).astype('int64')
                index.add_with_ids(vectors_np, ids_np)
                logger.info(f"向量添加完成。索引当前总向量数: {index.ntotal}")
            else:
                logger.warning("没有向量需要添加到 Faiss 索引。")

            if index.ntotal != len(id_map):
                 logger.error(f"严重错误：添加到 Faiss 的向量数 ({index.ntotal}) 与 ID 映射数 ({len(id_map)}) 不匹配！")
                 return False

            # --- 保存索引和映射文件 ---
            logger.info(f"准备保存 Faiss 索引到 {index_path}")
            faiss.write_index(index, index_path)

            logger.info(f"准备保存 ID 映射到 {map_path}")
            with open(map_path, 'wb') as f_map:
                pickle.dump(id_map, f_map)

            # --- 更新 IndexManager 的状态 ---
            self.semantic_index = index
            self.semantic_id_map = id_map
            self.semantic_index_enabled = True
            build_duration = time.time() - start_time
            logger.info(f"语义索引和映射构建并保存成功，耗时 {build_duration:.2f} 秒。")
            return True

        except faiss.FaissException as faiss_err:
            logger.exception(f"构建或保存 Faiss 索引时发生 Faiss 特定错误: {faiss_err}")
            self.semantic_index_enabled = False # 出错则禁用
            return False
        except Exception as e:
            logger.exception(f"构建或保存语义索引时发生未知错误: {e}")
            self.semantic_index_enabled = False
            return False


    # === 获取索引的接口 ===

    def get_minhash_lsh(self) -> Optional[Tuple[MinHashLSH, Dict[str, MinHash]]]:
        """获取 MinHash LSH 索引对象和语料库签名。"""
        if self.minhash_enabled:
            return self.lsh_index, self.corpus_minhashes
        else:
            return None

    def get_inverted_index(self) -> Optional[Dict[str, List[Tuple[str, int]]]]:
        """获取倒排索引字典。"""
        if self.inverted_index_enabled:
            return self.inverted_index
        else:
            return None

    def get_semantic_index(self) -> Optional[Tuple[faiss.Index, List[Tuple[str, int, int, int]]]]:
        """获取 Faiss 索引对象和 ID 映射。"""
        if self.semantic_index_enabled:
            return self.semantic_index, self.semantic_id_map
        else:
            return None

    def get_semantic_components(self) -> Optional[Tuple[faiss.Index, List[Tuple[str, int, int, int]]]]:
         """(兼容旧调用名) 获取 Faiss 索引对象和 ID 映射。"""
         return self.get_semantic_index()
